import cv2
import numpy as np
import mediapipe as mp
import time

# 初始化语义分割模型
segmentation = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=1)

# 初始化MediaPipe手部模型
mp_drawing = mp.solutions.drawing_utils
mp_hands = mp.solutions.hands
# 初始化手部检测器
hands = mp_hands.Hands(static_image_mode=False, max_num_hands=1, min_detection_confidence=0.5)

# 加载背景图像
new_size = (640, 480)
backgrounds_nums = 10
backgrounds = []
for i in range(0, backgrounds_nums):
    background_image = cv2.imread('backgrounds/'+str(i)+'.jpeg')
    background_image = cv2.resize(background_image, new_size)
    backgrounds.append(background_image)
current_background = -1     # -1表示不替换背景

# 打开摄像头
cap = cv2.VideoCapture(0)

# 避免背景连续变化
changeable = True
last_time = time.time()
delta = 0
change_cd = 1 # 1秒

while cap.isOpened():
    current_time = time.time()
    delta += current_time - last_time
    last_time = current_time
    if delta > change_cd:
        changeable = True
        delta = 0

    # 读取摄像头帧
    ret, frame = cap.read()

    # 修改大小
    frame = cv2.resize(frame, new_size)

    # 镜像翻转帧以匹配用户的视角
    frame = cv2.flip(frame, 1)
    
    # 将帧转换为RGB颜色空间
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    
    # 使用MediaPipe进行手部检测
    results = hands.process(frame_rgb)
    
    # 检测到手部
    if results.multi_hand_landmarks:
        for hand_landmarks in results.multi_hand_landmarks:
            # 获取手部关键点坐标
            landmarks = hand_landmarks.landmark
            
            # 获取左右手的关键点坐标, 判断向左指还是向右指
            if landmarks[mp_hands.HandLandmark.WRIST].x < landmarks[mp_hands.HandLandmark.INDEX_FINGER_TIP].x:
                direction = "right"
                if changeable:
                    changeable = False
                    if current_background == (backgrounds_nums - 1):
                        current_background = -1
                    else:
                        current_background += 1
            elif landmarks[mp_hands.HandLandmark.WRIST].x > landmarks[mp_hands.HandLandmark.INDEX_FINGER_TIP].x:
                direction = "left"
                if changeable:
                    changeable = False
                    if current_background < 0:
                        current_background = backgrounds_nums - 1
                    else:
                        current_background -= 1
            else:
                direction = "null"

            # 在帧上绘制手部关键点
            mp_drawing.draw_landmarks(frame, hand_landmarks, mp_hands.HAND_CONNECTIONS)
            # 在帧上显示识别结果
            cv2.putText(frame, direction, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    
    if current_background >= 0:
        # 将图像传递给MediaPipe进行语义分割
        image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = segmentation.process(image)

        # 提取语义分割掩码
        mask = np.stack((results.segmentation_mask,) * 3, axis=-1)

        # 将语义分割掩码应用到原始图像上
        background_image = backgrounds[current_background]
        frame = np.where(mask > 0, frame, background_image)

    # 显示帧
    cv2.imshow('Result', frame)
    
    # 按下Esc键退出
    if cv2.waitKey(1) == 27:
        break

# 释放摄像头并关闭窗口
cap.release()
cv2.destroyAllWindows()
